import rclpy
from rclpy.node import Node
from geometry_msgs.msg import Twist
from sensor_msgs.msg import LaserScan
from std_srvs.srv import Empty
from gymnasium import Env, spaces
import numpy as np

class TurtleBot3Env(Env):
    def __init__(self):
        super().__init__()

        # 初始化 ROS2 节点
        rclpy.init(args=None)
        self.node = rclpy.create_node("turtlebot3_env")

        # 发布动作 (cmd_vel)
        self.cmd_pub = self.node.create_publisher(Twist, "/cmd_vel", 10)

        # 订阅激光雷达
        self.scan_data = None
        self.scan_sub = self.node.create_subscription(
            LaserScan, "/scan", self.scan_callback, 10
        )

        # Gazebo reset services
        self.reset_sim_service_name = "/reset_simulation"
        self.reset_world_service_name = "/reset_world"
        self.reset_sim_client = self.node.create_client(Empty, self.reset_sim_service_name)
        self.reset_world_client = self.node.create_client(Empty, self.reset_world_service_name)

        # 动作空间：前进速度 + 旋转速度
        self.action_space = spaces.Box(
            low=np.array([0.0, -1.0]), high=np.array([0.3, 1.0]), dtype=np.float32
        )

        # 观测空间：雷达 360 度，简化成 24 个数
        self.observation_space = spaces.Box(
            low=0.0, high=3.5, shape=(24,), dtype=np.float32
        )

        self.done = False

    def scan_callback(self, msg: LaserScan):
        # 简化雷达：取均匀 24 个点
        ranges = np.array(msg.ranges)
        ranges = np.clip(ranges, 0, 3.5)
        step = len(ranges) // 24
        self.scan_data = ranges[::step][:24]

    def step(self, action):
        # 执行动作
        twist = Twist()
        twist.linear.x = float(action[0])
        twist.angular.z = float(action[1])
        self.cmd_pub.publish(twist)

        # 等待传感器数据
        rclpy.spin_once(self.node, timeout_sec=0.1)

        if self.scan_data is None:
            obs = np.ones(24) * 3.5
        else:
            obs = self.scan_data

        # ---- 奖励函数设计 ----
        min_dist = min(obs)
        reward = 0.0

        # 1) 前进奖励（速度越快越好）
        reward += twist.linear.x * 2.0

        # 2) 安全距离奖励（保持和障碍物 >0.5m 的奖励）
        if min_dist > 0.5:
            reward += 0.5
        else:
            reward -= 0.5

        # 3) 撞墙惩罚
        if min_dist < 0.2:
            reward = -5.0
            self.done = True

            # 撞墙后立刻停下
            self._stop_robot()
            self.node.get_logger().info("检测到碰撞，结束当前 episode")

        return obs, reward, self.done, False, {}

    def reset(self, seed=None, options=None):
        self.done = False
        # 发布停止命令，防止残余动作
        self._stop_robot()

        # 重置仿真环境
        self._reset_simulation()

        # 清理旧激光数据并等待新帧
        self.scan_data = None
        for _ in range(5):
            rclpy.spin_once(self.node, timeout_sec=0.1)

        obs = np.ones(24) * 3.5 if self.scan_data is None else self.scan_data
        return obs, {}

    def close(self):
        self.node.destroy_node()
        rclpy.shutdown()

    def _stop_robot(self):
        stop = Twist()
        stop.linear.x = 0.0
        stop.angular.z = 0.0
        self.cmd_pub.publish(stop)

    def _reset_simulation(self):
        request = Empty.Request()

        # 优先尝试 reset_world，失败则退回 reset_simulation
        if self._call_empty_service(self.reset_world_client, self.reset_world_service_name, request):
            return

        self._call_empty_service(self.reset_sim_client, self.reset_sim_service_name, request)

    def _call_empty_service(self, client, service_name, request, timeout_sec=2.0):
        if client is None:
            return False

        if not client.service_is_ready():
            if not client.wait_for_service(timeout_sec=timeout_sec):
                self.node.get_logger().warn(
                    f"服务 {service_name} 未准备就绪，跳过本次重置"
                )
                return False

        future = client.call_async(request)
        rclpy.spin_until_future_complete(self.node, future, timeout_sec=timeout_sec)

        if not future.done():
            self.node.get_logger().warn(
                f"调用服务 {service_name} 超时，可能需要手动检查 Gazebo"
            )
            return False

        if future.result() is None:
            self.node.get_logger().warn(
                f"服务 {service_name} 调用失败：{future.exception()}"
            )
            return False

        return True
